# plotting 3d data.
import numpy as np
import pandas as pd
import rrcf
from sklearn.ensemble import IsolationForest
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['font.size'] = 13
# plt.rcParams['font.sans-serif']=['SimHei']
# plt.rcParams['axes.unicode_minus'] = False
# mpl.rcParams['font.family'] = 'serif'
# mpl.rcParams['font.serif'] = 'Times New Roman'
# mpl.rcParams['font.style'] = 'normal'
# mpl.rcParams['font.variant'] = 'normal'
# mpl.rcParams['font.weight'] = 'normal'
# mpl.rcParams['font.stretch'] ='normal'
import time
# Read data
OilWell = pd.read_csv('TP15222.csv', index_col=0)
# OilWell = pd.read_csv('test_time.csv', index_col=0)
OilWell.index = pd.to_datetime(OilWell.index)
data_pre = OilWell['tubing'].astype(float).values
print(data_pre.shape)

#
data_pre1 = np.array(data_pre).reshape(2104, 3)
# data_pre1 = np.array(data_pre).reshape(100, 3)
#
data_pre1 = ((data_pre1 - data_pre1.min())
            / (data_pre1.max() - data_pre1.min()))

t1 = time.time()
# Run the RRCF algorithm and compute the CoDisp for each point.
# Set tree parameters
num_trees = 200
shingle_size = 49
tree_size = 300

# Use the "shingle" generator to create rolling window
points_pre = rrcf.shingle(data_pre, size = shingle_size)
points_pre = np.vstack([point for point in points_pre])
print(points_pre.shape)
n_pre = points_pre.shape[0]
sample_size_range_pre = (n_pre // tree_size, tree_size)

# Use the "shingle" generator to create rolling window
#
forest = []
while len(forest) < num_trees:
    ixs = np.random.choice(n_pre, size=sample_size_range_pre,
                           replace=False)
    trees = [rrcf.RCTree(points_pre[ix], index_labels=ix)
            for ix in ixs]
    forest.extend(trees)
    # print(len(forest))

t11 = time.time()
print("the time of construct rrcf :", t11-t1)
avg_codisp_pre = pd.Series(0.0, index=np.arange(n_pre))
index_pre = np.zeros(n_pre)

for tree in forest:
    # print(tree.leaves)
    # print(tree.root)
    codisp = pd.Series({leaf: tree.codisp(leaf) for leaf in tree.leaves})
    avg_codisp_pre[codisp.index] += codisp
    np.add.at(index_pre, codisp.index.values, 1)
print("the time of trace tree :", time.time()-t11)
avg_codisp_pre /= index_pre
avg_codisp_pre.index = OilWell.iloc[(shingle_size - 1):].index
t2 = time.time()
tt = t2 -t1
print("tt :", tt)
t1 = time.time()

# For comparison, we also run the Isolation Forest algorithm,
# contamination = OilWell['event'].sum()/len(OilWell)
contamination = 0.1
IF_pre = IsolationForest(n_estimators=num_trees, max_samples=500,
                         contamination=contamination,
                         behaviour='new',
                         random_state=0)
#
IF_pre.fit(points_pre)
if_scores_pre = IF_pre.score_samples(points_pre)

if_scores_pre = pd.Series(-if_scores_pre,index=(OilWell
                          .iloc[(shingle_size - 1):]
                          .index))
# Normalize anomaly scores to (0, 1)
avg_codisp_pre = ((avg_codisp_pre - avg_codisp_pre.min())
                  / (avg_codisp_pre.max() - avg_codisp_pre.min()))

if_scores_pre = ((if_scores_pre - if_scores_pre.min())
                  / (if_scores_pre.max() - if_scores_pre.min()))
#
threshold_pre = avg_codisp_pre.nlargest(n=10).min()
t2 = time.time()
tt = t2 -t1
print("tt :", tt)
#
x1 = (data_pre1[:2088, 0]-data_pre1[:2088, 0].max())+0.1
y1 = (data_pre1[:2088, 1]-data_pre1[:2088, 0].max())+0.1
z1 =  data_pre1[:2088, 2]-data_pre1[:2088, 0].min()
#
fig1 = plt.figure(figsize=(8, 6))
#
ax = fig1.add_subplot(111, projection='3d', label='RRCF')
sc1 = ax.scatter(data_pre1[:2088, 0], data_pre1[:2088, 1], data_pre1[:2088,2],
                linewidths=0.1, edgecolors='k',
                c = np.array(if_scores_pre).reshape(2088,3).max(axis = 1),
                cmap = 'jet')

sc2 = ax.scatter(x1, y1, z1,
                linewidths=0.1, edgecolors='k',
                c = np.array(avg_codisp_pre).reshape(2088,3).max(axis = 1),
                cmap = 'jet')

ax.set_xlim(-1.2, 1.2, 0.5)
ax.set_ylim(-1.2, 1.2, 0.5)
ax.set_xlabel('')
ax.set_ylabel('')
plt.title('Robust random cut forest vs Isolation forest', size = 14)
# Add a color bar which maps values to colors.
plt.colorbar(sc1, ax=ax, fraction=0.046, shrink=0.35, orientation='horizontal', pad=0.05)
            # cax=plt.axes([0.4, 0.08, 0.25, 0.01]))
#
# ax = fig1.add_subplot(122, projection='3d', label='RRCF')
# sc1 = ax.scatter(data_pre1[:2088, 0], data_pre1[:2088, 1], data_pre1[:2088,2],
#                 linewidths=0.1, edgecolors='k',
#                 c = np.array(if_scores_pre).reshape(2088,3).max(axis = 1),
#                 cmap = 'jet')
#
# sc2 = ax.scatter(x1, y1, z1,
#                 linewidths=0.1, edgecolors='k',
#                 c = np.array(if_scores_cur).reshape(2088,3).max(axis = 1),
#                 cmap = 'brg')
# ax.set_xlim(-1.0, 1.0)
# ax.set_ylim(-1.0, 1.0)
# #
# ax.set_xlabel('')
# ax.set_ylabel('')
# plt.title('Isolation forest')
# plt.tight_layout()
plt.tight_layout()
plt.savefig("fig_3.png", dpi=300)
plt.savefig("fig_3.svg", dpi=300)
plt.savefig("fig_3.pdf", dpi=300)
plt.savefig("fig_3.tif", dpi=300)

plt.show()


